"""图像处理核心模块"""
import cv2
import numpy as np
from typing import Tuple, Union

class ImageProcessor:
    """图像处理器"""
    
    @staticmethod
    def letterbox(img: np.ndarray, 
                new_shape: Union[int, Tuple[int, int]] = 640, 
                color: int = 114) -> Tuple[np.ndarray, float, Tuple[float, float]]:
        """
        图像缩放并添加边框
        
        Args:
            img: 输入图像
            new_shape: 目标尺寸
            color: 填充颜色
            
        Returns:
            处理后的图像, 缩放比例, 填充偏移量
        """
        shape = img.shape[:2]
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)
        
        # 计算缩放比例
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        
        # 缩放图像
        img_resized = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        
        # 计算填充
        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
        dw /= 2
        dh /= 2
        
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        
        # 根据图像通道数处理颜色值
        if len(img_resized.shape) == 3:  # 彩色图像
            if isinstance(color, int):
                fill_value = (color, color, color)
            else:
                fill_value = color
        else:  # 灰度图像
            fill_value = color if isinstance(color, int) else color[0]
        img_padded = cv2.copyMakeBorder(img_resized, 
                                        top, bottom, left, right, 
                                        cv2.BORDER_CONSTANT,
                                        value=(color,color,color))
        
        return img_padded, r, (dw, dh)
    
    @staticmethod
    def preprocess_image(img_path: str, 
                        imgsz: int = 640) -> Tuple[np.ndarray, np.ndarray, float, float, float]:
        """
        完整的图像预处理流程
        
        Args:
            img_path: 图像路径
            imgsz: 目标尺寸
            
        Returns:
            处理后的张量, 原始图像, 缩放比例, 水平偏移, 垂直偏移
        """
        img0 = cv2.imread(img_path)
        if img0 is None:
            raise ValueError(f"Cannot read image: {img_path}")
        
        img, r, (dw, dh) = ImageProcessor.letterbox(img0, new_shape=imgsz)
        
        # 转换为RGB并归一化
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None]  # NCHW
        
        return img, img0, r, dw, dh